# 新写法（0.2 之后推荐）
import logging

from langchain_community.document_loaders import (
    TextLoader,
    PyPDFLoader,
    Docx2txtLoader,
    UnstructuredExcelLoader,
)
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.vectorstores import Chroma

from langchain_huggingface import HuggingFaceEmbeddings
from langchain.storage import LocalFileStore
from langchain.embeddings.cache import CacheBackedEmbeddings
import hashlib
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(
    temperature=0,
    openai_api_key="sk-rgymzddvnnrgwawnkevkdzlxkxfvhbviqgibbfgpvsyoqtur",
    openai_api_base="https://api.siliconflow.cn/v1",
    model="deepseek-ai/DeepSeek-R1-0528-Qwen3-8B"
)

# 使用本地模型路径
local_model_path = r"D:\soft\sentence-transformers-master\bge-large-zh-v1.5"

embeddings = HuggingFaceEmbeddings(
    model_name=local_model_path,
    model_kwargs={'device': 'cpu'},
    encode_kwargs={'normalize_embeddings': True}
)


class ChatDoc():
    def __init__(self):
        self.doc = None
        self.splitText = []

    def getFile(self):
        doc = self.doc
        loaders = {
            "docx": Docx2txtLoader,
            "pdf": PyPDFLoader,
            "xlsx": UnstructuredExcelLoader
        }
        file_extension = doc.split(".")[-1]
        loader_class = loaders.get(file_extension)
        if loader_class:
            loader = loader_class(doc)
            text = loader.load()
            return text
        return None

    # 处理文档函数
    def splitSentences(self):
        full_text = self.getFile()
        if full_text is not None:
            text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=20)
            texts = text_splitter.split_documents(full_text)
            self.splitText = texts

    # 向量化与向量存储
    def embeddingAndVectorDB(self):
        db = Chroma.from_documents(documents=self.splitText, embedding=embeddings)
        return db

    # 提问并找到相关的文本块
    def askAndFindFiles(self, question):
        db = self.embeddingAndVectorDB()
        retriever = db.as_retriever()
        retriever_from_llm = MultiQueryRetriever.from_llm(retriever=retriever, llm=llm)
        unique_docs = retriever_from_llm.invoke(input=question)
        return unique_docs

chat_doc = ChatDoc()
chat_doc.doc = "file/若依环境使用手册.docx"
chat_doc.splitSentences()
# data = chat_doc.askAndFindFiles("若依基于Maven管理项目的构建，需要先安装好相应的版本。")
# print(data)


logging.basicConfig(level=logging.INFO)
logging.getLogger("langchain.retrievers.multi_query").setLevel(logging.DEBUG)
data = chat_doc.askAndFindFiles("怎么知道环境是否搭建成功")
print(data)
